﻿using ARPaintServer.Model.ObjectDetection;
using Yolov5Net.Scorer.Models;
using Yolov5Net.Scorer;
using SixLabors.ImageSharp;
using SixLabors.ImageSharp.PixelFormats;

namespace ARPaintServer.Services
{
	public class ObjectDetectionService:IObjectDetectionService
	{
		private const string yoloModelPath = "Asset/Model/yolov5s.onnx";

		public async Task<ItemResult> GetObjects(Model.ObjectDetection.ImageInfo imageInfo)
		{
			ItemResult itemResult;
			List<ItemInfo> imageItems = new();
			// 还原base64字节为图片流
			byte[] base64 = Convert.FromBase64String(imageInfo.ImageBase64);
			using (MemoryStream memoryStream = new(base64))
			{
				// 加载图片
				using var image = await Image.LoadAsync<Rgba32>(memoryStream);
				{
					// 加载图像检测模型
					using var scorer = new YoloScorer<YoloCocoP5Model>(yoloModelPath);
					{
						var predictions = scorer.Predict(image);
						foreach (var prediction in predictions) 
						{
							double score = Math.Round(prediction.Score, 2);
							Console.WriteLine($"发现 {prediction.Label.Name},位于{prediction.Rectangle.X},{prediction.Rectangle.Y}，置信度{score}");
							imageItems.Add(new()
							{
								Label = prediction.Label.Name,
								X = prediction.Rectangle.X,
								Y = prediction.Rectangle.Y,
								Height = prediction.Rectangle.Height,
								Width = prediction.Rectangle.Width,
								Score = prediction.Score
							});
						}
					}
				}
			}
			// 返回检测结果
			itemResult = new()
			{
				ImageName = imageInfo.ImageName,
				ImageItems = imageItems,
			};

			return itemResult;
		}
	}
}
